import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression

if __name__ == '__main__':
    path = 'data/Advertising.csv'

    # pandas读入
    data = pd.read_csv(path)  # TV、Radio、Newspaper、Sales
    x = data[['TV', 'Radio', 'Newspaper']]
    y = data['Sales']
    # 协方差矩阵,看特征与标签的关系
    print(data.corr())
    # 看散点图，,看特征与标签的关系
    plt.figure(facecolor='gray')
    plt.scatter(data['TV'], data['Sales'], s=30, edgecolor='white')
    plt.title('TV')
    plt.xlabel('TV')
    plt.ylabel('Sales')
    plt.show()

    x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.8, random_state=1)

    linreg = LinearRegression()
    model = linreg.fit(x_train, y_train)
    print(model)
    print(linreg.coef_, linreg.intercept_)

    y_hat = linreg.predict(x_test)
    mse = np.average((y_hat - np.array(y_test)) ** 2)  # Mean Squared Error
    rmse = np.sqrt(mse)  # Root Mean Squared Error
    print('MSE = ', mse, end=' ')
    print('RMSE = ', rmse)
    print('trainingR2 = ', linreg.score(x_train, y_train))
    print('testR2 = ', linreg.score(x_test, y_test))
